#!/usr/bin/env python3

"""
Lerobot Demo Launch - 启动使用 lerobot 硬件的完整 MoveIt 演示
"""

from ament_index_python.packages import get_package_share_directory
from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument, IncludeLaunchDescription, TimerAction
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch_ros.actions import Node
from launch_ros.substitutions import FindPackageShare
from moveit_configs_utils import MoveItConfigsBuilder
import os


def generate_launch_description():
    # Get package directories
    moveit_config_share = get_package_share_directory("scepter_moveit_config")
    lerobot_driver_share = get_package_share_directory("scepter_lerobot_driver")
    
    # Launch arguments
    robot_ip_arg = DeclareLaunchArgument(
        "robot_ip",
        default_value="localhost",
        description="IP address of the robot"
    )
    
    config_file_arg = DeclareLaunchArgument(
        "config_file",
        default_value="",
        description="Lerobot config file path"
    )
    
    rviz_config_arg = DeclareLaunchArgument(
        "rviz_config",
        default_value="moveit.rviz",
        description="RViz configuration file"
    )
    
    # Get launch configurations
    robot_ip = LaunchConfiguration("robot_ip")
    config_file = LaunchConfiguration("config_file")
    rviz_config = LaunchConfiguration("rviz_config")
    
    # Build MoveIt configuration for lerobot
    moveit_config = (
        MoveItConfigsBuilder("arm", package_name="scepter_moveit_config")
        .robot_description(
            file_path="config/lerobot.arm.urdf.xacro", 
            mappings={
                "robot_ip": robot_ip,
                "config_file": config_file
            }
        )
        .robot_description_semantic(file_path="config/arm.srdf")
        .planning_pipelines(pipelines=["ompl", "chomp", "pilz_industrial_motion_planner", "stomp"])
        .trajectory_execution(file_path="config/moveit_controllers.yaml")
        .planning_scene_monitor(
            publish_robot_description=True,
            publish_robot_description_semantic=True
        )
        .to_moveit_configs()
    )
    
    # 1. Lerobot Bridge Node
    lerobot_bridge_node = Node(
        package="scepter_lerobot_driver",
        executable="lerobot_bridge_node.py",
        name="lerobot_bridge_node",
        output="screen",
        parameters=[
            {"robot_config": "so101"},
            {"update_rate": 100.0},
            {"command_timeout": 1.0},
        ],
    )
    
    # 2. Robot State Publisher
    robot_state_publisher = Node(
        package="robot_state_publisher",
        executable="robot_state_publisher",
        name="robot_state_publisher",
        output="screen",
        parameters=[moveit_config.robot_description],
    )
    
    # 3. Static TF
    static_tf = Node(
        package="tf2_ros",
        executable="static_transform_publisher",
        name="static_transform_publisher",
        output="log",
        arguments=["--frame-id", "world", "--child-frame-id", "base_link"],
    )
    
    # 4. Controller Manager
    controller_manager = Node(
        package="controller_manager",
        executable="ros2_control_node",
        parameters=[
            moveit_config.robot_description,
            os.path.join(moveit_config_share, "config", "ros2_controllers.yaml"),
        ],
        output="screen",
    )
    
    # 5. MoveIt Move Group
    move_group_capabilities = {"capabilities": "move_group/ExecuteTaskSolutionCapability"}
    
    move_group_node = Node(
        package="moveit_ros_move_group",
        executable="move_group",
        output="screen",
        parameters=[
            moveit_config.to_dict(),
            move_group_capabilities,
        ],
    )
    
    # 6. Controller Spawners (with delays to ensure proper startup order)
    joint_state_broadcaster_spawner = TimerAction(
        period=2.0,
        actions=[
            Node(
                package="controller_manager",
                executable="spawner",
                arguments=["joint_state_broadcaster", "--controller-manager", "/controller_manager"],
                output="screen",
            )
        ],
    )
    
    arm_controller_spawner = TimerAction(
        period=4.0,
        actions=[
            Node(
                package="controller_manager",
                executable="spawner",
                arguments=["arm_group_controller", "-c", "/controller_manager"],
                output="screen",
            )
        ],
    )
    
    gripper_controller_spawner = TimerAction(
        period=6.0,
        actions=[
            Node(
                package="controller_manager",
                executable="spawner",
                arguments=["gripper_group_controller", "-c", "/controller_manager"],
                output="screen",
            )
        ],
    )
    
    # 7. RViz
    rviz_config_file = PathJoinSubstitution([
        FindPackageShare("scepter_moveit_config"),
        "config",
        rviz_config
    ])
    
    rviz_node = Node(
        package="rviz2",
        executable="rviz2",
        name="rviz2",
        output="log",
        arguments=["-d", rviz_config_file],
        parameters=[
            moveit_config.robot_description,
            moveit_config.robot_description_semantic,
            moveit_config.robot_description_kinematics,
            moveit_config.planning_pipelines,
            moveit_config.joint_limits,
        ],
    )
    
    return LaunchDescription([
        # Launch arguments
        robot_ip_arg,
        config_file_arg,
        rviz_config_arg,
        
        # Core nodes
        lerobot_bridge_node,
        robot_state_publisher,
        static_tf,
        controller_manager,
        move_group_node,
        
        # Controller spawners (with timing)
        joint_state_broadcaster_spawner,
        arm_controller_spawner,
        gripper_controller_spawner,
        
        # Visualization
        rviz_node,
    ])